//============================================================================
// Name        : Terra-GnR.cpp
// Author      : Aleksandar Stancic
//============================================================================

#include <iostream>
#include <cstdio>
#include <ctime>
#include <cstdlib>
#include <cmath>
#include <cstring>

#include "structs/math/util.h"
#include "structs/Bitmap.h"
#include "structs/PerlinSampler.h"
#include "structs/Camera.h"
#include "structs/FloatMap.h"
#include "structs/RGBMap.h"

#include "portutil/threading.h"

using namespace std;

#define WIDTH (640)
#define HEIGHT (360)

/// Parameters ///
//Vec3D sun(1.0f, 0.5f, 0.25f);
Vec3D sun(1.0f, 0.3f, 0.3f);
float normOffset = 0.1f;

int maxSteps = 4096;
//float waterLevel = -20.0f;
float waterLevel = -50.0f;

float rsMulti = 37.0f;
float rsInten = 0.02f;

//int shadowRays = 8;
int shadowRays = 32;
int AOrays = 1;
int waterRays = 16;
int refractRays = 8;

float ambientJitter = 1.0f;
//float shadowJitter = 0.05f;
float shadowJitter = 0.01f;
float waterJitter = 0.001f;
//float volumeJitter = 0.05f;
float volumeJitter = 0.01f;
float refractJitter = 0.003f;

float fogHeight = 0.8f;
float fogDensity = 1.0f;
float cFog = 0.2f;
float fogPower = 1.5f;

float ambientIntensity = 0.2f;
float waterFogDepth = 10.0f;
Vec3D baseWaterColor(0.8f * 128.0f * 0.1f, 0.9f * 128.0f * 0.1f, 1.0f * 128.0f * 0.1f);
/// Parameters ///

FloatMap heightMap(1024, 1024);
float terraMin = 0.0f;
float terraMax = 0.0f;

float step = 0.5f;

Vec3D dirtColor(167.0f / 255.0f, 124.0f / 255.0f, 63.0f / 255.0f);
//Vec3D stoneColor(163.0f / 255.0f, 168.0f / 255.0f, 163.0f / 255.0f);
Vec3D stoneColor(200.0f / 255.0f, 190.0f / 255.0f, 100.0f / 255.0f);
Vec3D waterStoneColor(63.0f / 255.0f, 49.0f / 255.0f, 25.0f / 255.0f);


struct RenderSpan
{
	int startX;
	int endX;
	int id;
	Bitmap *bmp;
};

float dotStart = 0.9f;
float dotEnd = 0.97f;

inline Vec3D getGroundColor(Vec3D normal, float h)
{
	Vec3D groundColor(52.0f / 255.0f, 72.0f / 255.0f, 14.0f / 255.0f);
	float dot = normal.dot(Vec3D::ZUnit);

	dot -= dotStart;
	dot = dot / (dotEnd - dotStart);
	clamp(dot, 0.0f, 1.0f);
	ease(dot);

	groundColor.lerpIP(dirtColor, 1.0f - dot);

	//h = abs(h);
	float submerged = -h;
	clampMin(h, 0.0f);
	submerged = submerged / 3.0f;
	clamp(submerged, 0.0f, 1.0f);
	h = 1.0f - h / 1.0f;

	h *= (dot * dot);
	clamp(h, 0.0f, 1.0f);

	groundColor.lerpIP(stoneColor, h);
	groundColor.lerpIP(waterStoneColor, submerged);

	return groundColor;
}

inline Vec3D getSkyColor(Vec3D ray)
{
	ray.normalizeIP();
	float sky = (float)abs(ray.dot(Vec3D::ZUnit));
	//sky *= sky;
	float r = LERP(212.0f, 90.0f, sky);
	float g = LERP(239.0f, 139.0f, sky);
	float b = LERP(224.0f, 169.0f, sky);

	float dot = ray.dot(sun);
	clampMin(dot, 0.0f);

	float sunDir = (dot - 0.95f) * 20.0f;
	clampMin(sunDir, 0.0f);
	sunDir = pow(sunDir, 500.0f) * 10.0f;
	clamp(sunDir, 0.0f, 1.0f);

	r += sunDir * 255.0f * 1.5f;
	g += sunDir * 255.0f * 1.5f;
	b += sunDir * 192.0f * 1.5f;


	float dotPow = pow(dot, 2.0f);
	dotPow -= 0.7f;
	clampMin(dotPow, 0.0f);
	dotPow *= 1.2f;
	clampMax(dotPow, 1.0f);

	r += dotPow * 255.0f * 0.5f;
	g += dotPow * 255.0f * 0.5f;
	b += dotPow * 192.0f * 0.5f;

	return Vec3D(r, g, b);
}

Vec3D rasterizeRay(Vec3D pos, Vec3D ray, float & Z);

Vec3D traceRay(Vec3D pos, Vec3D ray, int maxSteps, bool doBounces, bool doWater, float vecLen, float & length, float & volLit, float & volUnlit)
{
	//bool firstHit = false;
	float vlit = 0.0f;
	float vunl = 0.0f;
	float lastU = 0.0f;
	float lastL = 0.0f;
	int cooldown = 0;

	for(int i = 0; i < maxSteps; ++i)
	{
		pos += ray;
		length += vecLen;

		if(doBounces)
		{
			if(cooldown == 0)
			{
				float volSamp = 0.0f;
				Vec3D volRay = sun + Vec3D(randRangeF(-volumeJitter, volumeJitter),
						randRangeF(-volumeJitter, volumeJitter), randRangeF(-volumeJitter, volumeJitter));
				volRay.normalizeIP();
				traceRay(pos, volRay * 8.0f, 16, false, false, 4.0f, volSamp, vlit, vunl);

				float fogPow = (pos.Z - waterLevel) / fogHeight;
				clamp(fogPow, 0.0f, 1.0f);
				fogPow = pow(1.0f - fogPow, fogPower) * fogDensity;
				fogPow += cFog;

				if(volSamp == 0.0f)
				{
					//lastL = vecLen * pow((1.0f - ((pos.Z - waterLevel) / fogHeight)), fogPower);
					//clamp(lastL, 0.0f, vecLen);
					//lastL *= fogDensity;
					lastL = vecLen * fogPow; //fogDensity;
					lastU = 0.0f;
					volLit += lastL;
				}
				else
				{
					lastL = 0.0f;
					//lastU = vecLen;
					lastU = vecLen * fogPow; //pow((1.0f - ((pos.Z - waterLevel) / fogHeight)), fogPower);
					lastU *= fogDensity;
					clamp(lastU, 0.0f, vecLen);
					volUnlit += lastU;
				}
				cooldown = randRangeI(6, 10);
			}
			else
			{
				volLit += lastL;
				volUnlit += lastU;
				cooldown--;
			}
		}

		float samp = heightMap.reSample(pos.X, pos.Y, rsMulti, rsInten);

		if(pos.Z > terraMax && ray.Z > 0.0f)
		{
			length = 0.0f;
			return getSkyColor(ray);
		}

		if(samp > pos.Z || (doWater && doBounces && (pos.Z < waterLevel)))
		{
			//firstHit = true;
			pos -= ray;
			length -= vecLen;

			ray *= 0.1f;
			vecLen *= 0.1f;
			for(int n = 0; n < 12; ++n)
			{
				pos += ray;
				length += vecLen;
				samp = heightMap.reSample(pos.X, pos.Y, rsMulti, rsInten);

				if(doBounces && doWater && (pos.Z < waterLevel))
				{
					if(doBounces)
					{

						int occlusionCount = 0;


						for(int i = 0; i < shadowRays; ++i)
						{
							float shadowRay = 0.0f;
							Vec3D rayDir = sun;
							rayDir += Vec3D(randRangeF(-shadowJitter, shadowJitter), randRangeF(-shadowJitter, shadowJitter), randRangeF(-shadowJitter, shadowJitter));
							rayDir.normalizeIP();
							rayDir *= 2.0f;
							traceRay(Vec3D(pos.X, pos.Y, waterLevel + 0.001f), rayDir, 512, false, false, 2.0f, shadowRay, vlit, vunl);
							if(shadowRay > 0.0f)
								occlusionCount++;
						}
						float shadowTerm = ambientIntensity + (1.0f - (float)occlusionCount / shadowRays) * (1.0f - ambientIntensity);

						//float shadowTerm = 1.0f;

						//float ambientTerm = 1.0f;

						//float light = 1.0f;

						//float b = 0.25f + 0.25 * ambientTerm + light * shadowTerm * 0.5f;

						Vec3D waterColor = baseWaterColor;

						float waterLen = 0.0f;
						float waterV1 = 0.0f;
						float waterV2 = 0.0f;

						Vec3D refractColor;

						for(int i = 0; i < refractRays; ++i)
						{
							Vec3D refractRay = ray;
							refractRay += Vec3D(randRangeF(-refractJitter, refractJitter),
									randRangeF(-refractJitter, refractJitter), randRangeF(-refractJitter, refractJitter));
							refractRay.normalizeIP();
							refractRay *= 2.0f;

							Vec3D transpColor = traceRay(pos, refractRay, 32, true, false, 2.0f, waterLen, waterV1, waterV2);
							if(waterLen == 0.0f)
								waterLen = 1.0f;
							else
								waterLen = waterLen / (32 * 2.0f);
							//waterLen = pow(1.0f - waterLen, 2.0f);
							////waterLen = pow(waterLen, 3.0f);
							//waterLen = 1.0f - waterLen;
							////refractColor += transpColor.lerp(waterColor, waterLen);
							refractColor += transpColor.lerp(Vec3D::Zero, waterLen);
						}
						waterColor = refractColor / refractRays;

						Vec3D reflectedColor;
						for(int i = 0; i < waterRays; ++i)
						{
							float dist = 0.0f;
							//float aniso = randRangeF(0.0f, 1.0f);
							//aniso = pow(aniso, 15.0f);
							Vec3D reflRay = Vec3D(ray.X, ray.Y, abs(ray.Z));
							reflRay += Vec3D(randRangeF(-waterJitter, waterJitter), randRangeF(-waterJitter, waterJitter), randRangeF(-waterJitter, waterJitter)); //randRangeF(-waterJitter * 5.0f, waterJitter * 5.0f));
							reflRay.normalizeIP();
							//reflRay *= 10.0f;
							reflectedColor += rasterizeRay(Vec3D(pos.X, pos.Y, waterLevel + 0.01f), reflRay, dist);
						}
						reflectedColor /= waterRays;
						//waterColor *= shadowTerm;

						float fresnel = ray.normalize().dot(Vec3D::NegZUnit);
						fresnel = 0.5f + (pow(1.0f - abs(fresnel), 3.0f) * 0.5f);
						//waterLen *= waterLen;
						waterLen *= 10.0f;
						waterLen = pow(waterLen + 1.0f, 2.0f) - 1.0f;
						waterLen += 0.1f;
						clampMax(waterLen, 1.0f);
						//fresnel *= waterLen;
						return waterColor.lerp(reflectedColor, fresnel) * shadowTerm;
					}
					else
					{
						return Vec3D::One;
					}
				}
				else if(samp > pos.Z)
				{
					if(doBounces)
					{
						float sampX1 = heightMap.reSample(pos.X + normOffset, pos.Y, rsMulti, rsInten) - samp;
						float sampY1 = heightMap.reSample(pos.X, pos.Y + normOffset, rsMulti, rsInten) - samp;
						float sampX2 = heightMap.reSample(pos.X - normOffset, pos.Y, rsMulti, rsInten) - samp;
						float sampY2 = heightMap.reSample(pos.X, pos.Y - normOffset, rsMulti, rsInten) - samp;

						Vec3D nX((sampX1 - sampX2) * -0.5f, 0.0f, normOffset);
						Vec3D nY(0.0f, (sampY1 - sampY2) * -0.5f, normOffset);

						Vec3D normal = (nX + nY) * 0.5f;
						normal.normalizeIP();

						float light = normal.dot(sun);
						float normLight = light;
						if(light < 0.0f) light = 0.0f;
						light = 1.0f - pow(1.0f - light, 2.0f);
						light *= 1.7f;

						int occlusionCount = 0;

						for(int i = 0; i < shadowRays; ++i)
						{
							float shadowRay = 0.0f;
							Vec3D rayDir = sun;
							rayDir += Vec3D(randRangeF(-shadowJitter, shadowJitter), randRangeF(-shadowJitter, shadowJitter), randRangeF(-shadowJitter, shadowJitter));
							rayDir.normalizeIP();
							rayDir *= 2.0f;
							traceRay(Vec3D(pos.X, pos.Y, samp + 0.01f), rayDir, 512, false, false, 2.0f, shadowRay, vlit, vunl);
							if(shadowRay > 0.0f)
								occlusionCount++;
						}

						float shadowTerm = 1.0f - (float)occlusionCount / shadowRays;

						/*
						occlusionCount = 0;
						for(int i = 0; i < AOrays; ++i)
						{
							float aoRay = 0.0f;
							Vec3D rayDir = normal;
							//Vec3D rayDir = Vec3D::ZUnit;
							rayDir += Vec3D(randRangeF(-ambientJitter, ambientJitter), randRangeF(-ambientJitter, ambientJitter), randRangeF(-ambientJitter, ambientJitter));
							rayDir.normalizeIP();
							rayDir *= 0.05f;
							traceRay(Vec3D(pos.X, pos.Y, samp + 0.01f), rayDir, 64, false, false, 0.05f, aoRay, vlit, vunl);
							if(aoRay > 0.0f)
								occlusionCount++;
						}

						float ambientTerm = 0.5f + (1.0f - (float)occlusionCount / AOrays) * 0.5f;
*/
						float ambientTerm = 0.75f + normLight * 0.25f;

						float b = ambientIntensity * 255.0f * ambientTerm + light * shadowTerm * (1.0f - ambientIntensity) * 255.0f;
						//float b = 64 * ambientTerm + 32 * light * (1 - shadowTerm) + light * shadowTerm * 192.0f;

						Vec3D groundColor = getGroundColor(normal, pos.Z - waterLevel);


						if(pos.Z < waterLevel)
						{
							float depth = waterLevel - pos.Z;
							depth = depth / waterFogDepth;
							clamp(depth, 0.0f, 1.0f);
							depth = pow(1.0f - depth, 2.0f);
							depth = 1.0f - depth;
							//groundColor.lerpIP(baseWaterColor, depth);
							groundColor.lerpIP(Vec3D::Zero, depth);
						}

						return groundColor * b;
					}
					else
					{
						return Vec3D::One;
					}
				}
			}
			//printf("No detail-hit? %.2f %.2f %.2f", pos.X, pos.Y, pos.Z);
			break;
		}
	}

	length = 0.0f;
	return getSkyColor(ray);
}

Vec3D rasterizeRay(Vec3D pos, Vec3D ray, float & Z)
{
	ray *= step;

	float length = 0.0f;
	float volumeLit = 0.0f;
	float volumeUnlit = 0.0f;

	Vec3D color = traceRay(pos, ray, maxSteps, true, true, step, length, volumeLit, volumeUnlit);

	float ratioL = volumeLit / 512.0f;
	float ratioU = volumeUnlit / 512.0f;
	clampMax(ratioL, 1.0f);
	clampMax(ratioU, 1.0f);


	float dist = length / (step * maxSteps);

	if(dist > 0.0f)
	{
		color.lerpIP(Vec3D(212.0f * 0.2f, 239.0f * 0.2f, 224.0f * 0.2f), ratioU);
		color.lerpIP(Vec3D(212.0f, 239.0f, 224.0f), ratioL);
	}

	Z = length;
	return color;
}

float aspect = (float)WIDTH / (float)HEIGHT;

void *renderSpan(void * data)
{
	RenderSpan *spanData = (RenderSpan*)data;

	int sWIDTH = spanData->endX - spanData->startX;
	int offset = spanData->startX;

	Camera cam(aspect, 0.1f, 1000.0f);
	bool aaMark[HEIGHT][sWIDTH];
	memset(aaMark, 0, sWIDTH * HEIGHT * sizeof(bool));

	float aaSampX[4] = { 0.2f, -0.6f, -0.2f, 0.6f };
	float aaSampY[4] = { -0.6f, -0.2f, 0.6f, 0.2f };

	RGBMap rgbMap(sWIDTH, HEIGHT);
	FloatMap zBuffer(sWIDTH, HEIGHT);

	//cam.position.set(10.0f, 10.0f, 6.0f);
	//cam.target.set(30.0f, 30.0f, 3.0f);

	cam.position.set(380.0f, 120.0f, 60.0f);
	//cam.position.set(400.0f, 675.0f, 6.0f);
	cam.target.set(600.0f, 342.0f, -35.0f);

	cam.vertFOV = 50.0f;
	cam.up.set(0.0f, 0.0f, 1.0f);
	cam.update();

	for(Uint32 y = 0; y < HEIGHT; ++y)
	{
		for(Uint32 x = spanData->startX; x < spanData->endX; ++x)
		{
			float length = 0.0f;
			Vec3D ray = cam.getViewRay((float)x / (float)WIDTH - 0.5f, -(float)y / (float)HEIGHT + 0.5f);
			Vec3D color = rasterizeRay(cam.position, ray, length);

			rgbMap.pset(x - offset, y, color);
			zBuffer.pset(x - offset, y, length);
		}
		printf("[%d] Line %u / %u\n", spanData->id, y+1, HEIGHT);
	}

	float zThreshold = 10.0f;
	Vec3D redMarker(255.0f, 0.0f, 0.0f);

	for(Uint32 y = 0; y < HEIGHT - 1; ++y)
	{
		for(Uint32 x = spanData->startX; x < spanData->endX - 1; ++x)
		{
			float z0 = zBuffer.pget(x-offset, y);
			float z1 = zBuffer.pget(x+1-offset, y);
			float z2 = zBuffer.pget(x-offset, y+1);
			float z3 = zBuffer.pget(x+1-offset, y+1);

			if(abs(z0 - z1) > zThreshold)
			{
				//rgbMap.pset(x, y, redMarker);
				//rgbMap.pset(x+1, y, redMarker);
				aaMark[y][x-offset] = true;
				aaMark[y][x+1-offset] = true;
			}

			if(abs(z0 - z2) > zThreshold)
			{
				//rgbMap.pset(x, y, redMarker);
				//rgbMap.pset(x, y+1, redMarker);
				aaMark[y][x-offset] = true;
				aaMark[y+1][x-offset] = true;
			}

			if(abs(z0 - z3) > zThreshold)
			{
				//rgbMap.pset(x, y, redMarker);
				//rgbMap.pset(x+1, y+1, redMarker);
				aaMark[y][x-offset] = true;
				aaMark[y+1][x+1-offset] = true;
			}
		}
	}

	for(Uint32 y = 0; y < HEIGHT; ++y)
	{
		for(Uint32 x = spanData->startX; x < spanData->endX; ++x)
		{
			if(aaMark[y][x-offset])
			{
				float length = 0.0f;
				Vec3D color = rgbMap.pget(x-offset, y);
				for(int i = 0; i < 4; ++i)
				{
					Vec3D ray = cam.getViewRay((x + aaSampX[i]) / (float)WIDTH - 0.5f, -(y + aaSampY[i]) / (float)HEIGHT + 0.5f);
					color += rasterizeRay(cam.position, ray, length);
				}
				rgbMap.pset(x-offset, y, color / 5.0f);
			}
		}
	}

	for(Uint32 y = 0; y < HEIGHT; ++y)
	{
		for(Uint32 x = spanData->startX; x < spanData->endX; ++x)
		{
			//float z = zBuffer.pget(x, y);
			Vec3D color = rgbMap.pget(x-offset, y);

			clamp(color.X, 0.0f, 255.0f);
			clamp(color.Y, 0.0f, 255.0f);
			clamp(color.Z, 0.0f, 255.0f);

			spanData->bmp->pset(x-offset, y, (Uint32)color.X, (Uint32)color.Y, (Uint32)color.Z);
		}
	}

	//char fileName[64] = { 0 };
	//sprintf(fileName, "thread-%02d.bmp", spanData->id);
	//spanData->bmp->saveToBMP(fileName);

	return NULL;
}

int main() {
	Bitmap bmp(WIDTH, HEIGHT);
	PerlinSampler ps(0.5f, 9);

	int startTime = time(0);

	//ps.wrapAround = 1024;
	for(Uint32 y = 0; y < 1024; ++y)
	{
		float dy = ease(y / 1024.0f);
		for(Uint32 x = 0; x < 1024; ++x)
		{
			float dx = ease(x / 1024.0f);

			float r = (ps.multiSample(x, y) * (1.0f - dx) + ps.multiSample(1024 + x, y) * dx) * (1.0f - dy) +
					(ps.multiSample(x, 1024 + y) * dx + ps.multiSample(1024 + x, 1024 + y) * (1.0f - dx)) * (dy);
			//heightMap.pset(x, y, ps.multiSample(x, y) * 40.0f);
			heightMap.pset(x, y, r * 40.0f);
		}
	}
			//heightMap.pset(x, y, ps.sample(x, y, 0.0f, 0.0f) * 40.0f);

	terraMin = heightMap.getMin();
	terraMax = heightMap.getMax();

	fogHeight = (terraMax - waterLevel) * fogHeight;

	sun.normalizeIP();

	printf("Terragen done. Range: %.2f - %.2f\n", terraMin, terraMax);

	Bitmap heightMapOut(1024, 1024);
	for(Uint32 y = 0; y < 1024; ++y)
	{
		for(Uint32 x = 0; x < 1024; ++x)
		{
			Uint8 b = (Uint8)(((heightMap.pget(x, y) - terraMin) / (terraMax - terraMin)) * 255.0f);
			heightMapOut.pset(x, y, b, b, b, 255);
		}
	}
	heightMapOut.saveToBMP("heightmap.bmp");


	int threadCount = 4;

	RenderSpan* spans = new RenderSpan[threadCount];
	CUThread** threads = new CUThread*[threadCount];

	for(int i = 0; i < threadCount; ++i)
	{
		spans[i].id = i;
		spans[i].startX = (i * WIDTH) / threadCount;
		spans[i].endX = ((i + 1) * WIDTH) / threadCount + 1;
		spans[i].bmp = new Bitmap(spans[i].endX - spans[i].startX, HEIGHT);

		threads[i] = cu_spawnThread(&renderSpan, (void*)(&spans[i]));
	}

	for(int i = 0; i < threadCount; ++i)
		cu_joinThread(threads[i]);

	for(int i = 0; i < threadCount; ++i)
		spans[i].bmp->bltTo(&bmp, 0, 0, spans[i].startX, 0, spans[i].endX - spans[i].startX - 1, HEIGHT);
	bmp.saveToBMP("output.bmp");


	cout << "Done! " << (time(0) - startTime) << " secs" << endl;
	delete[] spans;
	delete[] threads;
	return 0;
}

